<?php

/**
 * This file is part of Webman AI.
 *
 * @author    walkor<walkor@workerman.net>
 * @copyright walkor<walkor@workerman.net>
 * @link      http://www.workerman.net/
 */

namespace plugin\ai\app\handler\driver;

use Workerman\Http\Client;
use Workerman\Http\Response;

class Gemini extends Base
{

    /**
     * 对话
     * @param $data
     * @param $options
     * @return void
     */
    public function completions($data, $options)
    {
        $headers = $this->getHeaders($options);
        if ($options['stream'] ?? false) {
            $data['stream'] = true;
        }
        $options = $this->formatOptions($options);
        $messages = static::formatMessages($data['messages']);
        $requestOptions = [
            'method' => 'POST',
            'data' => json_encode([
                'contents' => $messages
            ]),
            'headers' => $headers,
            'progress' => function ($buffer) use ($options) {
                static $tmp = '', $sentLen = 0;
                $tmp = rtrim($tmp . $buffer);
                if ($tmp[strlen($tmp) - 1] !== "]") {
                    $chunks = json_decode($tmp . ']', true);
                } else {
                    $chunks = json_decode($tmp, true);
                }
                if (!$chunks) {
                    return;
                }
                $removeLen = $sentLen;
                $sentLen = count($chunks);
                if ($sentLen > 0) {
                    array_splice($chunks, 0, $removeLen);
                }
                foreach ($chunks as $chunk) {
                    $chunk = array_merge(['content' => ''], $chunk);
                    $chunk['content'] = $chunk['candidates'][0]['content']['parts'][0]['text'] ?? '';
                    if ($chunk[0]['promptFeedback']['blockReason'] ?? false) {
                        $chunk['error'] = ['code' => 1, 'message' => '由于政策隐私等安全原因，此回复无法展示'];
                    }
                    $options['stream']($chunk);
                }
            },
            'success' => function (Response $response) use ($options) {
                $options['complete']($this->formatResponse((string)$response->getBody()), $response);
            },
            'error' => function ($exception) use ($options) {
                $options['complete']([
                    'error' => [
                        'code' => 'exception',
                        'message' => $exception->getMessage(),
                        'detail' => (string) $exception
                    ],
                ], new Response(0));
            }
        ];
        $model = $data['model'];
        $stream = $data['stream'] ?? false;
        $path = "/v1beta/models/$model:" . ($stream ? 'streamGenerateContent' : 'generateContent') . '?key=' . $this->apikey;
        $http = new Client(['timeout' => 600]);
        $http->request($this->api . $path, $requestOptions);
    }

    /**
     * @param $buffer
     * @return string
     */
    public static function formatResponse($buffer)
    {
        $data = json_decode($buffer, true);
        $message = '';
        if (is_array($data)) {
            foreach ($data as $chunk) {
                if (isset($chunk['error'])) {
                    $message .= $chunk['error']['message'] ?? '';
                } elseif (isset($chunk['candidates'][0]['content']['parts'][0]['text'])) {
                    $message .= $chunk['candidates'][0]['content']['parts'][0]['text'] ?? '';
                }
            }
        }
        return $message ?: $buffer;
    }

    /**
     * @param array $messages
     * @return array
     */
    public static function formatMessages(array $messages): array
    {
        $messages = parent::formatMessages($messages);
        foreach ($messages as $key => $value) {
            if (isset($value['parts'])) {
                break;
            }
            $value['parts'] = [['text' => $value['content']]];
            if (($value['role'] ?? '') == 'assistant') {
                $value['role'] = 'model';
            }
            unset($value['content']);
            $messages[$key] = $value;
        }
        return $messages;
    }
}